package com.aispace.supersql.example.controller;

import com.aispace.supersql.builder.TrainBuilder;
import com.aispace.supersql.engine.SpringSqlEngine;
import com.aispace.supersql.enumd.TrainPolicyType;
import com.aispace.supersql.vector.SpringVectorStore;
import com.alibaba.fastjson2.JSONObject;
import lombok.AllArgsConstructor;
import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.stream.Collectors;

@RestController
@RequestMapping("superSql")
@AllArgsConstructor
public class SuperSqlController{

    private final AzureOpenAiChatModel azureChatModel;

    private final SpringSqlEngine sqlEngine;

    private final SpringVectorStore store;

    @GetMapping("getSuperSql")
    public Object getSuperSql(@RequestParam String question) {
        String sql = sqlEngine.setChatModel(azureChatModel).generateSql(question);
        Object object = sqlEngine.executeSql(sql);
        return object;
    }



    @GetMapping("trainDdl")
    public String trainDDl() {
        String ddl = """
                 CREATE TABLE `dtp_hospital` (
                 	`id` BIGINT NOT NULL COMMENT '主键',
                 	`province` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '省份',
                 	`city` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '城市',
                 	`reporting_team` VARCHAR ( 50 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '提报团队',
                 	`district` VARCHAR ( 20 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '区',
                 	`hospital_name` VARCHAR ( 100 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '申请的DTP药房 主要对应的医院名称',
                 	`hospital_code` VARCHAR ( 50 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '申请DTP主要对应的医院code',
                 	`hospital_address` VARCHAR ( 255 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '医院具体地址',
                 	`location` VARCHAR ( 100 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '医院所在经纬度',
                 	`del_flag` INT NULL DEFAULT 0 COMMENT '删除状态 0正常 1已删除',
                 	`create_by` VARCHAR ( 32 ) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL COMMENT '创建人',
                 	`create_time` datetime NULL DEFAULT NULL COMMENT '创建时间',
                 	`update_by` VARCHAR ( 32 ) CHARACTER SET utf8 COLLATE utf8_general_ci NULL DEFAULT NULL COMMENT '更新人',
                 	`update_time` datetime NULL DEFAULT NULL COMMENT '更新时间',
                 	`image` VARCHAR ( 255 ) CHARACTER SET utf8 COLLATE utf8_unicode_ci NULL DEFAULT NULL COMMENT '图片',
                 	PRIMARY KEY ( `id` ) USING BTREE )
                 	ENGINE = INNODB CHARACTER SET = utf8 COLLATE = utf8_unicode_ci COMMENT = '医院表' ROW_FORMAT = Dynamic;
                """;
        sqlEngine.setChatModel(azureChatModel).train(TrainBuilder.builder().content(ddl).policy(TrainPolicyType.DDL).build());
        return "successful training";
    }


    @GetMapping("trainSql")
    public String trainSql() {
        String sql="SELECT * FROM DTP_HOSPITAL WHERE DISTRICT LIKE '%黄浦区%';";
        String question="在黄浦区的医院有哪些？";
        sqlEngine.setChatModel(azureChatModel).train(TrainBuilder.builder().content(sql).question(question).policy(TrainPolicyType.SQL).build());
        return "successful training";
    }


    @GetMapping("getSearchVector")
    public List<Document> getSearchVector(@RequestParam String question) {
        FilterExpressionBuilder expression = new FilterExpressionBuilder();
        List<Document> questionSqlList = this.store.searchByTag(
                question,
                expression.eq("script_type", TrainPolicyType.SQL.name()).build(),
                1);

        List<Document> collect = questionSqlList.stream()
                .filter(doc -> doc.getScore() >= 0.8)
                .collect(Collectors.toList());

        questionSqlList.forEach(document -> {
            JSONObject object = JSONObject.parseObject(document.getText());
            System.out.println(object.get("question"));
            System.out.println(object.get("sql"));
            System.out.println(document.getScore());

        });
        return collect;
    }


}
